data {
  int<lower=1> N; // length of data
  int<lower=1> k; // dimension of covariates
  vector[k] x[N];
  int<lower=0, upper=1> A[N];
  real<lower=0.01> length_scale_baseline; // length scale of Mater kernel
  real<lower=0.01> sigma_baseline; // sigma is the variance in matern kernel
  real<lower=0.01> length_scale_treatment; // length scale of Mater kernel in baseline effect
  real<lower=0.01> sigma_treatment; // sigma is the variance in matern kernel in treatment effect
  real y[N];
}
transformed data {
  real delta = 1e-9;
  matrix[N, N] L_baseline;
  matrix[N, N] L_treatment;
  matrix[N, N] K_baseline = gp_matern32_cov(x, sigma_baseline, length_scale_baseline);
  matrix[N, N] K_treatment = gp_matern32_cov(x, sigma_treatment, length_scale_treatment);
  // diagonal elements
  for (n in 1:N){
    K_baseline[n, n] = K_baseline[n, n] + delta;
    K_treatment[n, n] = K_treatment[n, n] + delta;
  }
  L_baseline = cholesky_decompose(K_baseline);
  L_treatment = cholesky_decompose(K_treatment);
}
parameters {
  vector[N] eta_baseline;
  vector[N] eta_treatment;
  real<lower=0.01> sigma;
}
transformed parameters{
  vector[N] f_baseline;
  vector[N] f_treatment;
  f_baseline = L_baseline * eta_baseline;
  f_treatment = L_treatment * eta_treatment;
}
model {
  eta_baseline ~ std_normal();
  eta_treatment ~ std_normal();
  sigma ~ inv_gamma(1,1);
  for (i in 1:N){
    y[i] ~ normal(f_baseline[i]+f_treatment[i]*A[i], sigma);
  }
}
generated quantities {
  real z0[N];
  real z1[N];
  for (i in 1:N){
    z0[i] = f_baseline[i];
    z1[i] = f_baseline[i]+f_treatment[i];
  }
}
